#include "ballistic_solver/ballistic_solver_core.hpp"
#include "rclcpp/rclcpp.hpp"
#include <limits>

namespace rm_aim
{
    double BallisticSolverCore::calcHorizontalDist(const geometry_msgs::msg::Point &target)
    {
        // x-y平面水平距离：sqrt(x² + y²)
        return sqrt(pow(target.x, 2) + pow(target.y, 2));
    }

    double BallisticSolverCore::calc3dDist(const geometry_msgs::msg::Point &target)
    {
        // 3D空间距离：sqrt(x² + y² + z²)
        return sqrt(pow(target.x, 2) + pow(target.y, 2) + pow(target.z, 2));
    }

    double BallisticSolverCore::solveYaw(const geometry_msgs::msg::Point &target)
    {
        // 纯几何Yaw角：atan2(y, x)（左为正，范围[-π, π]）
        return atan2(target.y, target.x);
    }

    double BallisticSolverCore::solvePitch(
        const geometry_msgs::msg::Point &target,
        double bullet_speed,
        double gravity,
        double k1)
    {
        double horizontal_dist = calcHorizontalDist(target);
        // 避免除零和无效计算
        if (horizontal_dist < 1e-6 || bullet_speed < 1e-6)
        {
            RCLCPP_DEBUG(rclcpp::get_logger("BallisticSolverCore"), 
                        "Invalid solvePitch input: horizontal_dist=%.6fm, bullet_speed=%.6fm/s",
                        horizontal_dist, bullet_speed);
            return 0.0;
        }

        // 初始几何角度（无重力/空气阻力补偿）
        double pitch_rad = atan2(target.z, horizontal_dist);

        // 迭代梯度下降补偿（最多100次迭代，误差<1e-6收敛）
        for (int iter = 0; iter < 100; ++iter)
        {
            double sin_theta = sin(pitch_rad);
            double cos_theta = cos(pitch_rad);

            // 避免cosθ为0导致除零
            if (fabs(cos_theta) < 1e-6)
            {
                RCLCPP_WARN(rclcpp::get_logger("BallisticSolverCore"), "cos(pitch) is too small, break iteration");
                break;
            }

            double sin2 = sin_theta * sin_theta;
            double cos2 = cos_theta * cos_theta;
            double cos3 = cos2 * cos_theta;

            // 计算当前Pitch角对应的飞行时间
            double flight_time = calcFlightTime(horizontal_dist, bullet_speed, pitch_rad, k1);

            // 计算实际弹道的z坐标（与目标z的偏差）
            double actual_z = (bullet_speed * sin_theta * flight_time) / cos_theta 
                           - (0.5 * gravity * pow(flight_time, 2)) / cos2;
            double delta_z = target.z - actual_z;

            // 收敛判断：z方向误差<1e-6m（约0.001mm）
            if (fabs(delta_z) < 1e-6)
            {
                break;
            }

            // 梯度下降修正Pitch角（导数计算）
            double derivate = (2 * gravity * flight_time * flight_time * sin_theta)/ cos3 
                            - (bullet_speed * flight_time * (1 + sin2))/ cos2;
            // 梯度修正
            pitch_rad -= delta_z / derivate;
        }

        return pitch_rad;
    }

    double BallisticSolverCore::calcFlightTime(
        double horizontal_dist,
        double bullet_speed,
        double pitch_rad,
        double k1)
    {
        double cos_theta = cos(pitch_rad);

        // 输入参数有效性检查
        if (fabs(cos_theta) < 1e-6 || bullet_speed < 1e-6 || horizontal_dist < 1e-6)
        {
            RCLCPP_WARN(rclcpp::get_logger("BallisticSolverCore"), 
                        "Invalid calcFlightTime input: cos_theta=%.6f, bullet_speed=%.6f, dist=%.6f",
                        cos_theta, bullet_speed, horizontal_dist);
            return 0.0;
        }

        // 空气阻力指数模型：t = (e^(k1*d) - 1) / (k1 * v0 * cosθ)
        return (exp(k1 * horizontal_dist) - 1) / (k1 * bullet_speed * cos_theta);
    }

    geometry_msgs::msg::Point BallisticSolverCore::predictArmorPos(
        const aim_interfaces::msg::Target &target,
        double dt)
    {
        geometry_msgs::msg::Point future_center_pos, armor_pos;
        double radius = target.radius_1;

        // 1. 预测目标中心未来位置（基于速度向量）
        future_center_pos.x = target.position.x + target.velocity.x * dt;
        future_center_pos.y = target.position.y + target.velocity.y * dt;
        future_center_pos.z = target.position.z + target.velocity.z * dt;
        double future_yaw = target.yaw + target.v_yaw * dt;

        // 计算装甲板相对于目标中心的位置
        armor_pos.x = future_center_pos.x - radius * cos(future_yaw);
        armor_pos.y = future_center_pos.y - radius * sin(future_yaw);
        armor_pos.z = future_center_pos.z;

        return armor_pos;
    }

    geometry_msgs::msg::Point BallisticSolverCore::calcPreciseAimPoint(
        const aim_interfaces::msg::Target &target,
        double total_delay,
        double flight_time_weight,
        double bullet_speed,
        double gravity,
        double k1)
    {
        geometry_msgs::msg::Point final_aim_point = target.position;

        // 初始目标水平距离（未预测前）
        double initial_horizontal_dist = calcHorizontalDist(final_aim_point);

        // 初始飞行时间（假设Pitch=0°，无补偿）
        double prev_fly_time = calcFlightTime(initial_horizontal_dist, bullet_speed, 0.0, k1);

        // 迭代求解：补偿总延迟+飞行时间（最多10次迭代）
        for (int iter = 0; iter < 10; ++iter)
        {
            // 预测时间 = 总延迟 + 飞行时间×权重
            double predict_dt = total_delay + (prev_fly_time * flight_time_weight);

            // 预测未来装甲板位置
            final_aim_point = predictArmorPos(target, predict_dt);

            // 计算预测位置的水平距离和对应的Pitch角
            double horizontal_dist = calcHorizontalDist(final_aim_point);
            double iterative_pitch = (horizontal_dist > 1e-6) 
                                   ? solvePitch(final_aim_point, bullet_speed, gravity, k1) 
                                   : 0.0;

            // 计算新的飞行时间
            double new_fly_time = calcFlightTime(horizontal_dist, bullet_speed, iterative_pitch, k1);

            // 收敛判断：飞行时间变化<0.001s（1ms）
            if (fabs(new_fly_time - prev_fly_time) < 0.001)
            {
                prev_fly_time = new_fly_time;
                break;
            }

            prev_fly_time = new_fly_time;
        }

        // 最终预测：基于收敛后的飞行时间
        double final_predict_dt = total_delay + prev_fly_time;
        final_aim_point = predictArmorPos(target, final_predict_dt);

        return final_aim_point;
    }
} // namespace rm_aim